package com.example.isoftlangchainai.rag.naive;

import com.example.isoftlangchainai.rag.utils.Utils;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.DocumentParser;
import dev.langchain4j.data.document.DocumentSplitter;
import dev.langchain4j.data.document.parser.apache.poi.ApachePoiDocumentParser;
import dev.langchain4j.data.document.splitter.DocumentSplitters;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.memory.ChatMemory;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.ChatModel;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.embedding.onnx.bgesmallenv15q.BgeSmallEnV15QuantizedEmbeddingModel;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import org.springframework.stereotype.Component;

import java.util.List;

import static dev.langchain4j.data.document.loader.FileSystemDocumentLoader.loadDocument;

/**
 * @Description: NaiveRagHelper
 * @Author :chenjun
 */
@Component
public class NaiveRagHelper {
    protected String baseUrl = "http://langchain4j.dev/demo/openai/v1";
    protected String apikey = "demo";
    protected String modelName = "gpt-4o-mini";

    public NaiveRagAssistant createAssistant(String documentPath) {

        System.out.println("1.let's create a chat model:聊天语言模型......");
        ChatModel chatModel = OpenAiChatModel.builder()
                .baseUrl(baseUrl)
                .apiKey(apikey)
                .modelName(modelName)
                .build();

        System.out.println("2.let's load a document that we want to use for RAG......");
        DocumentParser documentParser = new ApachePoiDocumentParser();
        //TextDocumentParser parser = new TextDocumentParser();
        //DocumentParser documentParser = new ApachePdfBoxDocumentParser();
        Document document = loadDocument(Utils.toPath(documentPath), documentParser);
//        Document document = ClassPathDocumentLoader.loadDocument("vectorfile/智慧城市.xlsx");

        System.out.println("3.split this document into smaller segments, also known as 'chunks'.....");
        //档分割策略：使用递归文档拆分器DocumentSplitter，优先按段落切分，若段落过长，则依次尝试按换行、句子、单词进行递归拆分，确保文本适配模型输入限制.
        DocumentSplitter splitter = DocumentSplitters.recursive(100, 0);
        List<TextSegment> segments = splitter.split(document);
        System.out.println("Total segments: " + segments.size());
        for (TextSegment segment : segments) {
            System.out.println("segments:" + segment.text());
        }
        System.out.println("4.create EmbeddingModel and embed (also known as 'vectorize') these segments......");
        //Embedding is needed for performing similarity searches. Langchain4j currently supports more than 10 popular embedding model providers
        EmbeddingModel embeddingModel = new BgeSmallEnV15QuantizedEmbeddingModel();
        List<Embedding> embeddings = embeddingModel.embedAll(segments).content();

        System.out.println("5.store these embeddings in an embedding store (also known as a 'vector database')......");
        //存储用于在每次与大语言模型（LLM）交互时，快速查找与用户查询最相关的文档片段,Langchain4j currently supports more than 15 popular embedding stores
        EmbeddingStore<TextSegment> embeddingStore = new InMemoryEmbeddingStore<>();
        embeddingStore.addAll(embeddings, segments);

        System.out.println("""
               6.创建内容检索器ContentRetriever，用于根据用户查询从向量化存储中检索最相关的文本片段。
                 配置返回相关性最高的结果为2、最低相似度得分为0.5......
                 """);
        ContentRetriever contentRetriever = EmbeddingStoreContentRetriever.builder()
                .embeddingStore(embeddingStore)
                .embeddingModel(embeddingModel)
                .maxResults(2)
                .minScore(0.7)
                .build();

        System.out.println("""
                7.使用聊天记忆（ChatMemory）来支持与大语言模型LLM的多轮对话，使其能够记住之前的交互内容......
                """);
        //按消息数量限制历史记录,如保留最近10条消息
        ChatMemory chatMemory = MessageWindowChatMemory.withMaxMessages(10);

        System.out.println("8.build our AI Service......");
        return AiServices.builder(NaiveRagAssistant.class)
                .chatModel(chatModel)//设置聊天语言模型
                .contentRetriever(contentRetriever)//设置内容检索器
                .chatMemory(chatMemory)//设置聊天记忆对象
                .build();
    }

}
